package com.yale.elasticsearch.utils;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.util.*;

/**
 * @author :yale
 * @date : 2024/7/31
 */
public class EsQueryUtil {
    private Map<String, Object> bool;
    private List<Map<String, Object>> must;
    private List<Map<String, Object>> sorts;
    private List<Map<String, Object>> ranges;
    private Integer from;
    private Integer size;
    private List<Object> after;

    public EsQueryUtil() {
        init();
    }
    public EsQueryUtil(int from, int size) {
        this.from = from;
        this.size = size;
        init();
    }

    private void init() {
        bool = new HashMap<>();
        must = new ArrayList<>();
        sorts = new ArrayList<>();
        ranges = new ArrayList<>();
        after = new ArrayList<>();
        bool.put("must", must);
    }

    /**
     * 添加精确查询条件
     * @param field 字段
     * @param values 值 set可防止条件重复
     */
    public void addTerms(String field, Set<String> values) {
        Map<String, Object> termClause = new HashMap<>();
        termClause.put("terms", new HashMap<String, Object>() {{
            put(field + ".keyword", values);
        }});
        must.add(termClause);
    }

    /**
     * 添加模糊查询条件
     * @param field 字段
     * @param value 值
     */
    public void addMatch(String field, String value) {
        Map<String, Object> matchClause = new HashMap<>();
        matchClause.put("match", new HashMap<String, Object>() {{
            put(field + ".keyword", value);
        }});
        must.add(matchClause);
    }

    /**
     * 添加排序条件
     * @param field 字段
     * @param order 值
     */
    public void addSort(String field, String order) {
        Map<String, Object> sortClause = new HashMap<>();
        sortClause.put(field, order.equals("asc") ? "asc" : "desc");
        sorts.add(sortClause);
    }

    /**
     * 添加范围条件
     * @param field 值
     * @param gt 开始
     * @param lt 结束
     */
    public void addRange(String field, Object gt, Object lt) {
        Map<String, Object> rangeClause = new HashMap<>();
        rangeClause.put("range", new HashMap<String, Object>() {{
            put(field, new HashMap<String, Object>() {{
                if (gt != null) {
                    put("gt", gt);
                }
                if (lt != null) {
                    put("lt", lt);
                }
            }});
        }});

        // 根据需求决定将范围查询添加到must、filter或其他bool子句中
        // 以下示例是添加到must中，根据实际情况调整
        ranges.add(rangeClause);
    }

    /**
     * 添加最后一条数据值，与排序字段顺序对应
     * @param data 值
     */
    public void addAfter(Object data) {
        after.add(data);
    }

    /**
     * 提取脚本
     * @param isCount 是否查询总数
     * @return 结果
     */
    public String getScript(boolean isCount) {
        // 确保terms和match条件位于bool的must子句中
        bool.put("must", must);

        // 将范围查询条件放入filter子句，如果存在的话
        if (!ranges.isEmpty()) {
            bool.put("filter", ranges);
        }

        Map<String, Object> query = new HashMap<>();
        query.put("bool", bool);

        Map<String, Object> finalQuery = new HashMap<>();
        if (!sorts.isEmpty()) {
            finalQuery.put("sort", sorts);
        }

        if (!isCount) {
            //  大于1w判断，如果大于1w，需要使用search_after(游标)进行分页
            if (from + size > 10000) {
                finalQuery.put("search_after", after);
            } else {
                finalQuery.put("from", from);
            }
            finalQuery.put("size", size);
        }
        finalQuery.put("query", query);
        finalQuery.put("track_total_hits", true);

        try {
            return new ObjectMapper().writeValueAsString(finalQuery);
        } catch (JsonProcessingException e) {
            return null;
        }
    }
}
